package com.warmheart.base.executor;

import java.util.Date;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.lang.NonNull;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import com.warmheart.base.config.IWhConfigService;
import com.warmheart.base.config.MyThreadPoolConfig;

public abstract class MdcThreadPoolTaskExecutor extends ThreadPoolTaskExecutor implements
        IWhConfigService<MyThreadPoolConfig> {

    private static Logger logger = LoggerFactory.getLogger(MdcThreadPoolTaskExecutor.class);

    private static final long serialVersionUID = 1L;

    final private boolean useFixedContext;

    final private Map<String, String> fixedContext;

    public static MdcThreadPoolTaskExecutor newWithInheritedMdc(MyThreadPoolConfig myThreadPoolConfig, TimeUnit unit,
            int queueCapacity) {
        return new MdcThreadPoolTaskExecutor(null, unit, queueCapacity) {

            private static final long serialVersionUID = 1L;

            @Override
            public MyThreadPoolConfig config() {
                return myThreadPoolConfig;
            }

        };
    }

    private MdcThreadPoolTaskExecutor(Map<String, String> fixedContext, TimeUnit unit, int queueCapacity) {
        setCorePoolSize(config().getCoreSize());
        setMaxPoolSize(config().getMaxSize());
        setKeepAliveSeconds((int) unit.toSeconds(config().getKeepAliveTime()));
        setQueueCapacity(queueCapacity);
        this.fixedContext = fixedContext;
        useFixedContext = (fixedContext != null);
    }

    private Map<String, String> getContextForTask() {
        return useFixedContext ? fixedContext : MDC.getCopyOfContextMap();
    }

    @Override
    public void execute(@NonNull Runnable command) {
        super.execute(wrap(command, getContextForTask()));
    }

    @NonNull
    @Override
    public Future<?> submit(@NonNull Runnable task) {
        return super.submit(wrap(task, getContextForTask()));
    }

    @NonNull
    @Override
    public <T> Future<T> submit(@NonNull Callable<T> task) {
        return super.submit(wrap(task, getContextForTask()));
    }

    private static <T> Callable<T> wrap(final Callable<T> task, final Map<String, String> context) {
        return () -> {
            Map<String, String> previous = MDC.getCopyOfContextMap();
            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            Date startTime = new Date();
            TraceContext.setSpan();
            String taskRemarks = "子线程执行任务";
            logger.info("----------------------------" + taskRemarks + "开始----------------------------");
            try {
                return task.call();
            } finally {
                Date endTime = new Date();
                logger.info("子线程处理耗时{}ms", endTime.getTime() - startTime.getTime());
                logger.info("----------------------------" + taskRemarks + "结束----------------------------");
                if (previous == null) {
                    MDC.clear();
                } else {
                    MDC.setContextMap(previous);
                }
            }
        };
    }

    private static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
        return () -> {
            Map<String, String> previous = MDC.getCopyOfContextMap();
            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            Date startTime = new Date();
            TraceContext.setSpan();
            String taskRemarks = "子线程执行任务";
            logger.info("----------------------------" + taskRemarks + "开始----------------------------");
            try {
                runnable.run();
            } finally {
                Date endTime = new Date();
                logger.info("子线程处理耗时{}ms", endTime.getTime() - startTime.getTime());
                logger.info("----------------------------" + taskRemarks + "结束----------------------------");
                if (previous == null) {
                    MDC.clear();
                } else {
                    MDC.setContextMap(previous);
                }
            }
        };
    }
}